/*
 * This code implements stVSSIM VQA as suggested by Anush Moorthy and A.C Bovik
 * in "Efficient Motion Weighted Spatio-Temporal Video SSIM Index"
 * http://live.ece.utexas.edu/publications/2010/moorthy_spie_jan10.pdf
 *
 * Please cite the following paper in any published work if you use this software.
 * Anus K. Moorthy, Alan C. Bovik,"Efficient Motion Weighted Spatio-Temporal
 * Video SSIM Index", SPIE Conference on Human Vision and Electronic Imaging,
 * San Jose, CA, January 17-21, 2010
 *
 * File: calcMV.cpp
 *
 * This Module Calculates the Motion vectors for the current frame using the
 * ARPS method.
 */

#include "vssim.h"

extern vssim_global_t global;

typedef struct
{
	CvPoint pt;			// Stores the Search point
	float cost;			// Stores the cost at the Search point
	bool valid; 		// If the cost is valid. Cost is used only if valid=true
}vssim_SearchPt_t;

/*
 * This function initializes the search point array
 */
static void initSearchPt(vssim_SearchPt_t *Pt, int n, int RoodSize)
{
	// center of the rood
	Pt[0].pt.x = 0; Pt[0].pt.y = 0;
	// Right end point of the rood
	Pt[1].pt.x = RoodSize; Pt[1].pt.y = 0;
	// Left end point of the rood
	Pt[2].pt.x = -RoodSize; Pt[2].pt.y = 0;
	// Top end point of the rood
	Pt[3].pt.x = 0; Pt[3].pt.y = RoodSize;
	// Bottom End point of the rood
	Pt[4].pt.x = 0; Pt[4].pt.y = -RoodSize;

	for (int i = 0; i < n; i++)
	{
		Pt[i].cost = 99999;				// > Max possible value of SAD (64*255)
		Pt[i].valid = false;
	}

	return;
}

/*
 * This funtion calculates the MAD score at each search point
 */
static float calcMAD(const unsigned char *CurrMB, const unsigned char *RefMB,
						int block_size, int stride)
{
	float MAD;
	float diff = 0.0;

	for (int i = 0; i < (block_size*block_size); i++)
	{
		diff += abs(CurrMB[i] - RefMB[i]);
	}
	MAD = diff/(float)(block_size*block_size);

	return MAD;
}

/*
 * This function determines which of the search points has the minimum cost
 */
static int findMINCost(vssim_SearchPt_t *SearchPts, int nSearch)
{
	int index = 0;
	float cost = 99999; //set to a max value>max possible value of the MB(255*64)

	for (int i=0; i < nSearch; i++)
	{
		if ((cost > SearchPts[i].cost) && (true == SearchPts[i].valid))
		{
			cost = SearchPts[i].cost;
			index = i;
		}
	}
	return index;
}

/*
 * This function returns the macro block of the requested size starting from
 * the co-ordinates passed.
 */
static void getMB(const int w, const int h, int stride,
					const unsigned char *Fr, int block_size,
					unsigned char *MB)
{
	int index =0;
	for (int j = h; j < h+block_size; j++)
	{
		for (int i = w; i < w+block_size; i++)
		{
			MB[index] = Fr[j*stride + i];
			index++;
		}
	}

	return;
}

/*
 * This function calculates the motion vector for every 16th frame by using the
 * previous frame as its reference
 */
extern void calcMotionVectors( FILE *pfile, int cntframe,
									vssim_size_t size, vssim_ImgComp_t comp,
									vssim_mv_t *mv, int block_size, int num_blocks)
{
	// allocate memory for the frames
	unsigned char *CurrFrame = (unsigned char *)malloc(sizeof(unsigned char)*size.width*size.height);
	unsigned char *PrevFrame = (unsigned char *)malloc(sizeof(unsigned char)*size.width*size.height);

	if(CurrFrame == NULL || PrevFrame == NULL)
	{
		printf("Memory allocation for current frame/previous frame failed. Exiting.. \n");
		exit(-1);
	}

	unsigned char CurrMB[64];
	unsigned char RefMB[64];

	long offset;
	switch (comp)
	{
		case Y:
			offset = 0;
			break;
		case U:
			offset = global.size.width * HEIGHT_Y(&global);
			break;
		case V:
			offset = global.size.width *(HEIGHT_Y(&global)+HEIGHT_U(&global));
			break;
	}

	/* Read Y/U/V component of the current frame in the file */
	fseek(pfile, ((cntframe*global.size.width*FRAME_HEIGHT(&global))+offset), SEEK_SET);
	fread(CurrFrame, sizeof(char), (size.width * size.height), pfile);

	/* Read the Y/U/V component of the prev frame from the file */
	fseek(pfile, (((cntframe-1)*global.size.width*FRAME_HEIGHT(&global))+offset), SEEK_SET);
	fread(PrevFrame, sizeof(char), (size.width * size.height), pfile);

	int stride = size.width;

	// Maintains the count of the num of blocks
	int nblock = 0;
	// size of the Rood ARM, default = 2 assuming it is the left most block
	int RoodSize = 2;
	// There are at the most 6 search points in ARPS, assuming default = 5
	int nSearch = 5;
	// stores the search points for the array
	vssim_SearchPt_t *SearchPts;
	SearchPts = (vssim_SearchPt_t *)malloc(6 * sizeof(vssim_SearchPt_t));

	memset(mv, 0, (sizeof(vssim_mv_t))*num_blocks);

	for (int h = 0; h < (size.height - block_size + 1); h+=block_size)
	{
		for (int w = 0; w < size.width; w+=block_size)
		{
			bool MatchFound = false;

			/* Get the MB from the Curr frame for which the MV is to be calculated */
			getMB(w, h, stride, CurrFrame, block_size, CurrMB);

			/* If the Left most MB is under consideration the RoodSize = 2
			 * else the RoodSize is decided from the Prev Blocks MV */
			if (w == 0 || nblock == 0)
			{
				RoodSize = 2;
				nSearch = 5;
			}
			else
			{
				/* If not the left most block then MV is initialized to prev block
				 * MV  and Rood size = max of the X/Y */
				mv[nblock].x = mv[nblock-1].x;
				mv[nblock].y = mv[nblock-1].y;
				RoodSize = MAX(abs(mv[nblock].x), abs(mv[nblock].y));

				/* Check if the MV is along one of the rood arms, if not
				 * it creates the 6th search point */
				if ( ((abs(mv[nblock].x) != 0) && (abs(mv[nblock].y) == 0)) ||
					 ((abs(mv[nblock].x) == 0) && (abs(mv[nblock].y) != 0))
				    )
				{
					nSearch = 5;
				}
				else
				{
					nSearch = 6;
				}
			}
			// If the rood size = 0, then only the center point is my point of search
			if (RoodSize == 0)
			{
				nSearch = 1;
			}

			/* Create and populate the Search Points Array with the co-ordinates
			 * 5/6 points that are to be searched for */
			initSearchPt(SearchPts, nSearch, RoodSize);

			if (nSearch == 6)
			{
				// Motion vector from the previous
				SearchPts[5].pt.x = mv[nblock].x; SearchPts[5].pt.y = mv[nblock].y;
			}

			/* Get the MB at each point from the Previous frame and calculate the
			 * MAD with the MB from the Current frame */
			// Calculate the cost at each MB and store the cost in the array
			for (int i=0; i< nSearch; i++)
			{
				//move to the search point
				int x = w;
				x += SearchPts[i].pt.x;

				int y = h;
				y += SearchPts[i].pt.y;

				// Check the boundary conditions
				if ((x < 0) || (x+block_size > size.width-1) ||
					(y < 0) || (y+block_size > size.height-1))
				{
					SearchPts[i].cost = -1;
					SearchPts[i].valid = false;
					continue;	// out of the boundary
				}

				// Get the corresponding MB from the Ref plane
				getMB(x, y, stride, PrevFrame, block_size, RefMB);

				// Calculate the cost
				SearchPts[i].cost = calcMAD(CurrMB, RefMB, block_size, stride);
				SearchPts[i].valid = true;
			}

			// Find the block with Min Cost
			int index = findMINCost(SearchPts, nSearch);

			//Store the location of min MAD
			CvPoint NewSearch;
			NewSearch.x = w + SearchPts[index].pt.x;
			NewSearch.y = h + SearchPts[index].pt.y;

			RoodSize = 1;	// Use the Unit size of Rood patter for finer search
			nSearch = 5;	// Do not have to search the center point as we
							// already have the score for it

			vssim_SearchPt_t URPSearch[5];
			initSearchPt(URPSearch, nSearch, RoodSize);

			URPSearch[0].cost = SearchPts[index].cost;	// copy the cost from the
													// prev large Rood pattern search
			URPSearch[0].valid = true;

			while (false == MatchFound)
			{
				// start from 1, 0- center point we already have the values
				for (int i=1; i<nSearch; i++)
				{
					// move to the search points
					int x = NewSearch.x;
					x += URPSearch[i].pt.x;

					int y = NewSearch.y;
					y += URPSearch[i].pt.y;

					// Check the boundary conditions
					if ((x < 0) || (x+block_size > size.width-1) ||
						(y < 0) || (y+block_size > size.height-1))
					{
						URPSearch[i].cost = -1;
						URPSearch[i].valid = false;
						continue;	// out of the boundary
					}

					// Get the corresponding MB from the Ref plane
					getMB(x, y, stride, PrevFrame, block_size, RefMB);

					// Calculate the cost
					URPSearch[i].cost = calcMAD(CurrMB, RefMB, block_size, stride);
					URPSearch[i].valid = true;
				}

				// Find the block with Min Cost
				int index = findMINCost(URPSearch, nSearch);

				if (index == 0)
				{
					// update the motion vector if a match is found
					MatchFound = true;
					mv[nblock].x = NewSearch.x - w;
					mv[nblock].y = NewSearch.y - h;
				}
				else
				{
					NewSearch.x += URPSearch[index].pt.x;
					NewSearch.y += URPSearch[index].pt.y;
					URPSearch[0].cost = URPSearch[index].cost;	// copy the cost from the
																// prev large Rood pattern search
					URPSearch[0].valid = true;
				}
			}
			nblock++;
		}
	}

	/* clean up */
	free(SearchPts);
	free(CurrFrame);
	free(PrevFrame);

	return;
}
